import pickle
from absl import app
from absl import flags
import matplotlib

matplotlib.use('Agg')  # 无图形界面模式
import matplotlib.pyplot as plt
import numpy as np
import os
from pyevtk.hl import pointsToVTK
from matplotlib import animation
from matplotlib.patches import Circle
from mpl_toolkits.mplot3d.art3d import Line3DCollection

# 原始参数
flags.DEFINE_string("rollout_dir", None, help="Directory where rollout.pkl are located")
flags.DEFINE_string("rollout_name", None, help="Name of rollout `.pkl` file")
flags.DEFINE_integer("step_stride", 100, help="Stride of steps to skip.")
flags.DEFINE_bool("change_yz", False, help="Change y and z axis.")
flags.DEFINE_enum("output_mode", "gif", ["gif", "vtk"], help="Type of render output")
# 新增：特定时间步截图参数
flags.DEFINE_string("snapshot_steps", "0,96,196,296,396,496,596", help="Comma-separated list of steps to save as PNG snapshots")

FLAGS = flags.FLAGS

# 粒子物理属性
PARTICLE_PROPERTIES = {
    6: {"density": 2680.0, "radius": 0.002},  # 沙: 密度(kg/m³), 半径(m)
    0: {"density": 2196.0, "radius": 0.004},  # 固体: 密度(kg/m³), 半径(m)
}

TYPE_TO_COLOR = {
    1: "red",  # for droplet
    3: "black",  # Boundary particles.
    0: "green",  # Rigid solids.
    7: "magenta",  # Goop.
    6: "gold",  # Sand.
    5: "blue",  # Water.
}

# 滚筒参数 - 根据用户提供的图片信息调整
DRUM_CENTER = np.array([0.015, 0.055, 0.055])  # 中心位置
DRUM_INNER_RADIUS = 0.05  # 内径
DRUM_OUTER_RADIUS = 0.055  # 外径
DRUM_LENGTH = 0.03  # 长度（X方向）


# ==== 修改后的Render类 ====
class Render():
    def __init__(self, input_dir, input_name):
        # 解析截图时间步
        self.snapshot_steps = list(map(int, FLAGS.snapshot_steps.split(','))) if FLAGS.snapshot_steps else []
        self.input_dir = input_dir
        self.input_name = input_name
        self.output_dir = input_dir
        self.output_name = input_name

        with open(f"{self.input_dir}{self.input_name}.pkl", "rb") as file:
            rollout_data = pickle.load(file)
        self.rollout_data = rollout_data

        rollout_cases = [
            ["ground_truth_rollout", "Reality"],
            ["predicted_rollout", "GNS"]
        ]
        self.rollout_cases = rollout_cases

        trajectory = {}
        for rollout_case in rollout_cases:
            trajectory[rollout_case[0]] = np.concatenate(
                [rollout_data["initial_positions"], rollout_data[rollout_case[0]]], axis=0
            )
        self.trajectory = trajectory
        self.loss = self.rollout_data['loss'].item() if 'loss' in self.rollout_data else 0.0

        # 轨迹信息
        self.dims = trajectory[rollout_cases[0][0]].shape[2]
        self.num_particles = trajectory[rollout_cases[0][0]].shape[1]
        self.num_steps = trajectory[rollout_cases[0][0]].shape[0]
        self.boundaries = rollout_data["metadata"]["bounds"]
        self.particle_type = rollout_data["particle_types"]

        # 计算粒子大小（基于半径）
        self.particle_sizes = np.zeros(self.num_particles)
        for i, p_type in enumerate(self.particle_type):
            if p_type in PARTICLE_PROPERTIES:
                self.particle_sizes[i] = PARTICLE_PROPERTIES[p_type]["radius"] * 20000  # 转换为适合绘制的尺寸
            else:
                self.particle_sizes[i] = 1  # 默认大小


    def color_map(self):
        color_map = np.empty(self.num_particles, dtype="object")
        for material_id, color in TYPE_TO_COLOR.items():
            color_index = np.where(np.array(self.particle_type) == material_id)
            color_map[color_index] = color
        return list(color_map)


    def color_mask(self):
        color_mask = []
        for material_id, color in TYPE_TO_COLOR.items():
            mask = np.array(self.particle_type) == material_id
            if mask.any():
                color_mask.append([mask, color])
        return color_mask

    def draw_drum(self, ax, change_yz=False):
        """绘制滚筒"""
        if self.dims != 3:
            return  # 只在3D中绘制滚筒

        # 生成滚筒的点
        theta = np.linspace(0, 2 * np.pi, 100)

        # 滚筒的X方向范围（向前）
        x_start = DRUM_CENTER[0] - DRUM_LENGTH / 2
        x_end = DRUM_CENTER[0] + DRUM_LENGTH / 2

        # 内圆柱 - 在YZ平面（向右和向上）
        y_inner = DRUM_INNER_RADIUS * np.cos(theta) + DRUM_CENTER[1]
        z_inner = DRUM_INNER_RADIUS * np.sin(theta) + DRUM_CENTER[2]

        # 外圆柱 - 在YZ平面
        y_outer = DRUM_OUTER_RADIUS * np.cos(theta) + DRUM_CENTER[1]
        z_outer = DRUM_OUTER_RADIUS * np.sin(theta) + DRUM_CENTER[2]

        # 填充圆筒区域为天蓝色（半透明）
        if not change_yz:
            # 创建圆筒表面的网格
            from mpl_toolkits.mplot3d.art3d import Poly3DCollection

            # 创建外圆柱面
            X_outer = np.array([x_start * np.ones_like(y_outer), x_end * np.ones_like(y_outer)])
            Y_outer = np.array([y_outer, y_outer])
            Z_outer = np.array([z_outer, z_outer])

            # 创建内圆柱面
            X_inner = np.array([x_start * np.ones_like(y_inner), x_end * np.ones_like(y_inner)])
            Y_inner = np.array([y_inner, y_inner])
            Z_inner = np.array([z_inner, z_inner])

            # 创建前后端面
            # 前端面（x_start）
            verts_front = []
            for i in range(len(theta) - 1):
                verts_front.append([(x_start, y_outer[i], z_outer[i]),
                                    (x_start, y_outer[i + 1], z_outer[i + 1]),
                                    (x_start, y_inner[i + 1], z_inner[i + 1]),
                                    (x_start, y_inner[i], z_inner[i])])

            # 后端面（x_end）
            verts_back = []
            for i in range(len(theta) - 1):
                verts_back.append([(x_end, y_outer[i], z_outer[i]),
                                   (x_end, y_outer[i + 1], z_outer[i + 1]),
                                   (x_end, y_inner[i + 1], z_inner[i + 1]),
                                   (x_end, y_inner[i], z_inner[i])])

            # 创建并添加多边形集合
            poly_front = Poly3DCollection(verts_front, alpha=0.3, linewidths=0, edgecolors='none')
            poly_front.set_facecolor('skyblue')
            ax.add_collection3d(poly_front)

            poly_back = Poly3DCollection(verts_back, alpha=0.3, linewidths=0, edgecolors='none')
            poly_back.set_facecolor('skyblue')
            ax.add_collection3d(poly_back)

            # 创建并添加圆柱面
            ax.plot_surface(X_outer, Y_outer, Z_outer, color='skyblue', alpha=0.3)
            ax.plot_surface(X_inner, Y_inner, Z_inner, color='skyblue', alpha=0.3)

            # 绘制框架轮廓
            ax.plot(x_start, y_inner, z_inner, 'gray', linewidth=1, alpha=0.7)
            ax.plot(x_start, y_outer, z_outer, 'gray', linewidth=1, alpha=0.7)
            ax.plot(x_end, y_inner, z_inner, 'gray', linewidth=1, alpha=0.7)
            ax.plot(x_end, y_outer, z_outer, 'gray', linewidth=1, alpha=0.7)

            # 绘制连接线
            for i in range(0, 100, 10):
                ax.plot([x_start, x_start],
                        [y_inner[i], y_outer[i]],
                        [z_inner[i], z_outer[i]], 'gray', linewidth=1, alpha=0.7)
                ax.plot([x_end, x_end],
                        [y_inner[i], y_outer[i]],
                        [z_inner[i], z_outer[i]], 'gray', linewidth=1, alpha=0.7)
        else:
            # 交换y和z轴的情况
            # 创建圆筒表面的网格
            from mpl_toolkits.mplot3d.art3d import Poly3DCollection

            # 创建外圆柱面
            Y_outer = np.array([x_start * np.ones_like(y_outer), x_end * np.ones_like(y_outer)])
            X_outer = np.array([y_outer, y_outer])
            Z_outer = np.array([z_outer, z_outer])

            # 创建内圆柱面
            Y_inner = np.array([x_start * np.ones_like(y_inner), x_end * np.ones_like(y_inner)])
            X_inner = np.array([y_inner, y_inner])
            Z_inner = np.array([z_inner, z_inner])

            # 创建前后端面
            # 前端面（x_start）
            verts_front = []
            for i in range(len(theta) - 1):
                verts_front.append([(y_outer[i], x_start, z_outer[i]),
                                    (y_outer[i + 1], x_start, z_outer[i + 1]),
                                    (y_inner[i + 1], x_start, z_inner[i + 1]),
                                    (y_inner[i], x_start, z_inner[i])])

            # 后端面（x_end）
            verts_back = []
            for i in range(len(theta) - 1):
                verts_back.append([(y_outer[i], x_end, z_outer[i]),
                                   (y_outer[i + 1], x_end, z_outer[i + 1]),
                                   (y_inner[i + 1], x_end, z_inner[i + 1]),
                                   (y_inner[i], x_end, z_inner[i])])

            # 创建并添加多边形集合
            poly_front = Poly3DCollection(verts_front, alpha=0.3, linewidths=0, edgecolors='none')
            poly_front.set_facecolor('skyblue')
            ax.add_collection3d(poly_front)

            poly_back = Poly3DCollection(verts_back, alpha=0.3, linewidths=0, edgecolors='none')
            poly_back.set_facecolor('skyblue')
            ax.add_collection3d(poly_back)

            # 创建并添加圆柱面
            ax.plot_surface(X_outer, Y_outer, Z_outer, color='skyblue', alpha=0.3)
            ax.plot_surface(X_inner, Y_inner, Z_inner, color='skyblue', alpha=0.3)

            # 绘制框架轮廓
            ax.plot(y_inner, x_start, z_inner, 'gray', linewidth=1, alpha=0.7)
            ax.plot(y_outer, x_start, z_outer, 'gray', linewidth=1, alpha=0.7)
            ax.plot(y_inner, x_end, z_inner, 'gray', linewidth=1, alpha=0.7)
            ax.plot(y_outer, x_end, z_outer, 'gray', linewidth=1, alpha=0.7)

            # 绘制连接线
            for i in range(0, 100, 10):
                ax.plot([y_inner[i], y_outer[i]],
                        [x_start, x_start],
                        [z_inner[i], z_outer[i]], 'gray', linewidth=1, alpha=0.7)
                ax.plot([y_inner[i], y_outer[i]],
                        [x_end, x_end],
                        [z_inner[i], z_outer[i]], 'gray', linewidth=1, alpha=0.7)


    # def render_snapshot(self, step, positions, case_name):
    #     """渲染并保存特定时间步的PNG截图"""
    #     # 创建输出目录
    #     snapshot_dir = os.path.join(self.output_dir, "snapshots")
    #     os.makedirs(snapshot_dir, exist_ok=True)
    #
    #     # 创建图形
    #     fig = plt.figure(figsize=(10, 8))
    #
    #     if self.dims == 2:
    #         ax = fig.add_subplot(111)
    #     elif self.dims == 3:
    #         ax = fig.add_subplot(111, projection='3d')
    #
    #     # 设置边界 - 根据用户提供的图片信息调整
    #     xboundary = [0, 0.03]  # x维度0-0.03
    #     yboundary = [0, 0.11]  # y维度0-0.11
    #     zboundary = [0, 0.11]  # z维度0-0.11
    #
    #     color_mask = self.color_mask()
    #
    #     # 绘制滚筒
    #     if self.dims == 3:
    #         self.draw_drum(ax, FLAGS.change_yz)
    #
    #     # 绘制
    #     if self.dims == 2:
    #         ax.set_aspect("equal")
    #         ax.set_xlim([float(xboundary[0]), float(xboundary[1])])
    #         ax.set_ylim([float(yboundary[0]), float(yboundary[1])])
    #         for mask, color in color_mask:
    #             sizes = self.particle_sizes[mask]
    #             ax.scatter(positions[mask, 0], positions[mask, 1], s=sizes, color=color)
    #         ax.grid(True, which='both')
    #         ax.set_title(f"{case_name} at step {step}")
    #
    #     elif self.dims == 3:
    #         if not FLAGS.change_yz:
    #             ax.set_xlim([float(xboundary[0]), float(xboundary[1])])
    #             ax.set_ylim([float(yboundary[0]), float(yboundary[1])])
    #             ax.set_zlim([float(zboundary[0]), float(zboundary[1])])
    #             for mask, color in color_mask:
    #                 sizes = self.particle_sizes[mask]
    #                 ax.scatter(positions[mask, 0], positions[mask, 1], positions[mask, 2], s=sizes, color=color)
    #             ax.set_box_aspect((float(xboundary[1]) - float(xboundary[0]),
    #                                float(yboundary[1]) - float(yboundary[0]),
    #                                float(zboundary[1]) - float(zboundary[0])))
    #             ax.view_init(elev=20, azim=0)  # 固定视角
    #         else:
    #             ax.set_xlim([float(xboundary[0]), float(xboundary[1])])
    #             ax.set_yl极im([float(zboundary[0]), float(zboundary[1])])
    #             ax.set_zlim([float(yboundary[0]), float(yboundary[1])])
    #             for mask, color in color_mask:
    #                 sizes = self.particle_sizes[mask]
    #                 ax.scatter(positions[mask, 0], positions[mask, 2], positions[mask, 1], s=sizes, color=color)
    #             ax.set_box_aspect((float(xboundary[1]) - float(xboundary[0]),
    #                                float(zboundary[1]) - float(zboundary[0]),
    #                                float(yboundary[1]) - float(yboundary[0])))
    #             ax.view_init(elev=20, azim=0)
    #         ax.grid(True, which='both')
    #         ax.set_title(f"{case_name} at step {step}")
    #
    #     # 保存图像
    #     plt.tight_layout()
    #     plt.savefig(os.path.join(snapshot_dir, f"{self.output_name}_{case_name}_step{step}.png"), dpi=300)
    #     plt.close(fig)
    #     print(f"截图已保存: {os.path.join(snapshot_dir, f'{self.output_name}_{case_name}_step{step}.png')}")
    def render_snapshot(self, step, positions, case_name):
        """渲染并保存特定时间步的PNG截图"""
        # 创建输出目录
        snapshot_dir = os.path.join(self.output_dir, "snapshots")
        os.makedirs(snapshot_dir, exist_ok=True)

        # 创建图形
        fig = plt.figure(figsize=(10, 8))

        if self.dims == 2:
            ax = fig.add_subplot(111)
        elif self.dims == 3:
            ax = fig.add_subplot(111, projection='3d')

        # 设置边界 - 根据用户提供的图片信息调整
        xboundary = [0, 0.03]  # x维度0-0.03
        yboundary = [0, 0.11]  # y维度0-0.11
        zboundary = [0, 0.11]  # z维度0-0.11

        color_mask = self.color_mask()

        # 绘制滚筒
        if self.dims == 3:
            self.draw_drum(ax, FLAGS.change_yz)

        # 绘制
        if self.dims == 2:
            ax.set_aspect("equal")
            ax.set_xlim([float(xboundary[0]), float(xboundary[1])])
            ax.set_ylim([float(yboundary[0]), float(yboundary[1])])

            # 移除坐标轴和刻度
            ax.set_axis_off()

            for mask, color in color_mask:
                sizes = self.particle_sizes[mask]
                ax.scatter(positions[mask, 0], positions[mask, 1], s=sizes, color=color)
            ax.grid(False)  # 移除网格
            ax.set_title(f"{case_name} at step {step}")

        elif self.dims == 3:
            if not FLAGS.change_yz:
                ax.set_xlim([float(xboundary[0]), float(xboundary[1])])
                ax.set_ylim([float(yboundary[0]), float(yboundary[1])])
                ax.set_zlim([float(zboundary[0]), float(zboundary[1])])

                # 移除坐标轴和刻度
                ax.set_axis_off()

                for mask, color in color_mask:
                    sizes = self.particle_sizes[mask]
                    ax.scatter(positions[mask, 0], positions[mask, 1], positions[mask, 2], s=sizes, color=color)
                ax.set_box_aspect((float(xboundary[1]) - float(xboundary[0]),
                                   float(yboundary[1]) - float(yboundary[0]),
                                   float(zboundary[1]) - float(zboundary[0])))
                ax.view_init(elev=20, azim=0)  # 固定视角
            else:
                ax.set_x极lim([float(xboundary[0]), float(xboundary[1])])
                ax.set_ylim([float(zboundary[0]), float(zboundary[1])])
                ax.set_zlim([float(yboundary[0]), float(yboundary[1])])

                # 移除坐标轴和刻度
                ax.set_axis_off()

                for mask, color in color_mask:
                    sizes = self.particle_sizes[mask]
                    ax.scatter(positions[mask, 0], positions[mask, 2], positions[mask, 1], s=sizes, color=color)
                ax.set_box_aspect((float(xboundary[1]) - float(xboundary[0]),
                                   float(zboundary[1]) - float(zboundary[0]),
                                   float(yboundary[1]) - float(yboundary[0])))
                ax.view_init(elev=20, azim=0)
            ax.grid(False)  # 移除网格
            ax.set_title(f"{case_name} at step {step}")

        # 保存图像
        plt.tight_layout()
        plt.savefig(os.path.join(snapshot_dir, f"{self.output_name}_{case_name}_step{step}.png"), dpi=300)
        plt.close(fig)
        print(f"截图已保存: {os.path.join(snapshot_dir, f'{self.output_name}_{case_name}_step{step}.png')}")

    def render_gif_animation(self, point_size=1, timestep_stride=3,
                             vertical_camera_angle=20, viewpoint_rotation=0.5,
                             change_yz=False):
        fig = plt.figure()
        if self.dims == 2:
            ax1 = fig.add_subplot(1, 2, 1, projection='rectilinear')
            ax2 = fig.add_subplot(1, 2, 2, projection='rectilinear')
            axes = [ax1, ax2]
        elif self.dims == 3:
            ax1 = fig.add_subplot(1, 2, 1, projection='3d')
            ax2 = fig.add_subplot(1, 2, 2, projection='3d')
            axes = [ax1, ax2]

        trajectory_datacases = [self.rollout_cases[0][0], self.rollout_cases[1][0]]
        render_datacases = [self.rollout_cases[0][1], self.rollout_cases[1][1]]

        xboundary = self.boundaries[0]
        yboundary = self.boundaries[1]
        if self.dims == 3:
            zboundary = self.boundaries[2]

        color_mask = self.color_mask()

        def animate(i):
            print(f"Render step {i}/{self.num_steps}")

            for j, datacase in enumerate(trajectory_datacases):
                positions = self.trajectory[datacase][i]

                # 检查是否为需要截图的步数
                if i in self.snapshot_steps:
                    self.render_snapshot(i, positions, render_datacases[j])

                if self.dims == 2:
                    axes[j] = fig.add_subplot(1, 2, j + 1)
                    axes[j].set_aspect("equal")
                    axes[j].set_xlim([float(xboundary[0]), float(xboundary[1])])
                    axes[j].set_ylim([float(yboundary[0]), float(yboundary[1])])
                    for mask, color in color_mask:
                        axes[j].scatter(positions[mask, 0], positions[mask, 1],
                                        s=point_size, color=color)
                    axes[j].grid(True, which='both')
                    axes[j].set_title(render_datacases[j])

                elif self.dims == 3:
                    # 绘制滚筒
                    self.draw_drum(axes[j], change_yz)

                    axes[j] = fig.add_subplot(1, 2, j + 1, projection='3d')
                    if not change_yz:
                        axes[j].set_xlim([float(xboundary[0]), float(xboundary[1])])
                        axes[j].set_ylim([float(yboundary[0]), float(yboundary[1])])
                        axes[j].set_zlim([float(zboundary[0]), float(zboundary[1])])
                        for mask, color in color_mask:
                            axes[j].scatter(positions[mask, 0], positions[mask, 1],
                                            positions[mask, 2], s=point_size, color=color)
                        axes[j].set_box_aspect((float(xboundary[1]) - float(xboundary[0]),
                                                float(yboundary[1]) - float(yboundary[0]),
                                                float(zboundary[1]) - float(zboundary[0])))
                        axes[j].view_init(elev=vertical_camera_angle, azim=i * viewpoint_rotation)

                    else:
                        axes[j].set_xlim([float(xboundary[0]), float(xboundary[1])])
                        axes[j].set_ylim([float(zboundary[0]), float(zboundary[1])])
                        axes[j].set_zlim([float(yboundary[0]), float(yboundary[1])])
                        for mask, color in color_mask:
                            axes[j].scatter(positions[mask, 0], positions[mask, 2],
                                            positions[mask, 1], s=point_size, color=color)
                        axes[j].set_box_aspect((float(xboundary[1]) - float(xboundary[0]),
                                                float(zboundary[1]) - float(zboundary[0]),
                                                float(yboundary[1]) - float(yboundary[0])))
                        axes[j].view_init(elev=vertical_camera_angle, azim=i * viewpoint_rotation)

                    axes[j].grid(True, which='both')
                    axes[j].set_title(render_datacases[j])

            fig.suptitle(f"{i}/{self.num_steps}, Total MSE: {self.loss:.2e}")

        ani = animation.FuncAnimation(
            fig, animate, frames=np.arange(96, self.num_steps, timestep_stride), interval=50)

        ani.save(f'{self.output_dir}{self.output_name}.gif', dpi=300, fps=10, writer='imagemagick')
        print(f"Animation saved to: {self.output_dir}{self.output_name}.gif")


    def write_vtk(self):
        for rollout_case, label in self.rollout_cases:
            path = f"{self.output_dir}{self.output_name}_vtk-{label}"
            if not os.path.exists(path):
                os.makedirs(path)

            initial_position = self.trajectory[rollout_case][0]

            for i, coord in enumerate(self.trajectory[rollout_case]):
                # 检查是否为需要截图的步数
                if i in self.snapshot_steps:
                    self.render_snapshot(i, coord, label)

                # 确保位移计算使用副本而不是视图
                disp = np.linalg.norm(coord - initial_position, axis=1)

                # 处理2D情况
                if self.dims == 2:
                    z_data = np.zeros(coord.shape[0], dtype=np.float32)
                else:
                    z_data = coord[:, 2].astype(np.float32)

                # 确保所有数组都是C连续且类型正确
                x_data = np.ascontiguousarray(coord[:, 0].astype(np.float32))
                y_data = np.ascontiguousarray(coord[:, 1].astype(np.float32))
                z_data = np.ascontiguousarray(z_data)
                disp_data = np.ascontiguousarray(disp.astype(np.float32))

                pointsToVTK(
                    f"{path}/points{i}",
                    x_data,
                    y_data,
                    z_data,
                    data={"displacement": disp_data}
                )

        print(f"vtk saved to: {self.output_dir}{self.output_name}...")


# ==== 主函数 ====
def main(_):
    if not FLAGS.rollout_dir:
        raise ValueError("A `rollout_dir` must be passed.")
    if not FLAGS.rollout_name:
        raise ValueError("A `rollout_name`must be passed.")

    render = Render(input_dir=FLAGS.rollout_dir, input_name=FLAGS.rollout_name)

    if FLAGS.output_mode == "gif":
        render.render_gif_animation(
            point_size=1,
            timestep_stride=FLAGS.step_stride,
            vertical_camera_angle=20,
            viewpoint_rotation=0.0,
            change_yz=FLAGS.change_yz
        )
    elif FLAGS.output_mode == "vtk":
        render.write_vtk()


if __name__ == '__main__':
    app.run(main)